# scepter_motion_visualized.py - 带 RViz 可视化的版本
import time
import os

# 设置环境变量
os.environ['ROS_PACKAGE_PATH'] = '/home/parallels/PycharmProjects/aghanim-scepter/src:/opt/ros/jazzy/share'

from moveit.task_constructor import core, stages
import rclcpp


def main():
    rclcpp.init()
    node = rclcpp.Node("scepter_mtc_visualized")
    
    print("🎬 Starting Scepter MTC Visualization Demo...")
    print("📺 Make sure RViz is running to see the visualization!")
    print("   Run: ros2 launch scepter_moveit_config demo.launch.py")
    print()
    
    # 设置规划组
    arm_group = "arm_group"
    gripper_group = "gripper_group"

    # 创建任务并加载机器人模型
    task = core.Task()
    task.name = "scepter_visualization_demo"
    
    try:
        task.loadRobotModel(node)
        print("✅ Robot model loaded successfully")
    except Exception as e:
        print(f"❌ Failed to load robot model: {e}")
        return False

    # 创建关节插值规划器
    joint_planner = core.JointInterpolationPlanner()
    print("✅ Joint space planner created")
    
    # 从当前状态开始
    current_state = stages.CurrentState("current state")
    task.add(current_state)

    print("📋 Building visualization-friendly task sequence...")
    print("⏱️  Each stage will be executed with pauses for observation\n")

    # 任务序列：演示完整的机械臂运动
    stages_info = []
    
    # 1. 移动到准备位置
    print("🏗️  Stage 1: Moving to ready position...")
    move_to_ready = stages.MoveTo("move to ready position", joint_planner)
    move_to_ready.group = arm_group
    move_to_ready.setGoal("ready_pose")
    task.add(move_to_ready)
    stages_info.append("Move to ready position")

    # 2. 打开夹爪
    print("🏗️  Stage 2: Opening gripper...")
    open_gripper = stages.MoveTo("open gripper", joint_planner)
    open_gripper.group = gripper_group
    open_gripper.setGoal("open_pose")
    task.add(open_gripper)
    stages_info.append("Open gripper")

    # 3. 移动到完成位置
    print("🏗️  Stage 3: Moving to finish position...")
    move_to_finish = stages.MoveTo("move to finish position", joint_planner)
    move_to_finish.group = arm_group
    move_to_finish.setGoal("finish_pose")
    task.add(move_to_finish)
    stages_info.append("Move to finish position")

    # 4. 关闭夹爪
    print("🏗️  Stage 4: Closing gripper...")
    close_gripper = stages.MoveTo("close gripper", joint_planner)
    close_gripper.group = gripper_group
    close_gripper.setGoal("close_pose")
    task.add(close_gripper)
    stages_info.append("Close gripper")

    # 5. 回到准备位置
    print("🏗️  Stage 5: Moving back to ready...")
    move_back_ready = stages.MoveTo("move back to ready", joint_planner)
    move_back_ready.group = arm_group
    move_back_ready.setGoal("ready_pose")
    task.add(move_back_ready)
    stages_info.append("Move back to ready position")

    # 6. 回到初始位置
    print("🏗️  Stage 6: Returning home...")
    move_home = stages.MoveTo("move home", joint_planner)
    move_home.group = arm_group
    move_home.setGoal("stand_pose")
    task.add(move_home)
    stages_info.append("Return to home position")

    print(f"\n📊 Task created with {len(stages_info)} stages")
    
    # 执行规划
    print("\n🔄 Planning complete trajectory...")
    input("👆 Press ENTER when RViz is ready to continue...")
    
    try:
        planning_success = task.plan()
        
        if planning_success:
            print("\n✅ Planning succeeded!")
            print("🎯 Trajectory will be visualized in RViz")
            
            # 检查解决方案
            if hasattr(task, 'solutions') and len(task.solutions) > 0:
                solution = task.solutions[0]
                print(f"📈 Found solution with trajectory")
                
                print("\n🎬 Starting execution visualization...")
                print("📺 Watch RViz to see the robot movement!")
                print("⏳ The trajectory will be published and should animate in RViz\n")
                
                # 发布解决方案到 RViz
                task.publish(solution)
                print("✅ Trajectory published to RViz!")
                
                # 显示阶段信息并等待
                print("\n📋 Execution stages:")
                for i, stage_name in enumerate(stages_info):
                    print(f"  {i+1}. 🎯 {stage_name}")
                    if i < len(stages_info) - 1:
                        time.sleep(2)  # 在每个阶段之间暂停
                
                print("\n✨ All stages completed!")
                print("🔄 The animation should loop in RViz")
                
                # 保持运行以便观察
                print("\n⏳ Keeping the trajectory alive for observation...")
                print("💡 You can see the planned path in RViz Motion Planning display")
                print("🔄 The planned path should be looping for visualization")
                print("🛑 Press Ctrl+C to stop when you're done observing\n")
                
                try:
                    while True:
                        time.sleep(1)
                        # 继续发布以保持可视化
                        task.publish(solution)
                        time.sleep(4)  # 每5秒重新发布一次
                except KeyboardInterrupt:
                    print("\n🛑 Visualization stopped by user")
                
                return True
            else:
                print("⚠️ Planning succeeded but no solutions found")
                return False
        else:
            print("❌ Planning failed")
            print("💡 Check RViz console for error details")
            return False
            
    except Exception as e:
        print(f"❌ Error during planning: {e}")
        return False


def print_rviz_instructions():
    """打印 RViz 设置说明"""
    print("\n" + "="*60)
    print("📺 RViz 可视化设置说明")
    print("="*60)
    print("1. 🚀 启动 RViz（在另一个终端中运行）：")
    print("   cd /home/parallels/PycharmProjects/aghanim-scepter")
    print("   ros2 launch scepter_moveit_config demo.launch.py")
    print()
    print("2. 📋 在 RViz 中检查以下显示项：")
    print("   ✅ MotionPlanning - 显示机器人模型和轨迹")
    print("   ✅ Planned Path - 显示规划的路径（应该是绿色）")
    print("   ✅ Loop Animation - 确保开启以看到动画循环")
    print()
    print("3. 🎛️ 调整 RViz 设置：")
    print("   - State Display Time: 0.05s (控制动画速度)")
    print("   - Loop Animation: True (循环播放)")
    print("   - Show Trail: True (显示轨迹痕迹)")
    print()
    print("4. 👀 观察要点：")
    print("   - 机器人模型应该按照规划路径移动")
    print("   - 每个关节的运动应该平滑")
    print("   - 夹爪的开合动作")
    print("="*60)


if __name__ == "__main__":
    try:
        print_rviz_instructions()
        print("\n⏳ Waiting 3 seconds for you to read the instructions...")
        time.sleep(3)
        
        success = main()
        
        if success:
            print("\n🎉 Visualization demo completed successfully!")
        else:
            print("\n⚠️ Demo completed with issues")
            
    except KeyboardInterrupt:
        print("\n🛑 Demo interrupted by user")
    except Exception as e:
        print(f"\n💥 Unexpected error: {e}")
        print("🔧 Please check your RViz setup and try again")